package amazon;

/*
 * myList[] may be null, in that case, the function should return null.
 * If k is negative or zero, return null.
 * Assume that (x * x + y * y) of each Point stays within the range of data type DOUBLE.
 *
 * My approach:
 * My initial thought was to sort the whole array in place using the sort function in java.util.Arrays,
 * based on the square of the distance from a Point to the origin point based on the fact that 
 * distance * distance = (x * x + y * y),
 * and return the first k elements from the sorted array.
 * This takes an average time of O(n lg(n)), where n is the length of the array.
 *
 * To do better, I would create a Max-Heap of size k from the first k elements from the arrray, then every time
 * I add another element into the heap, I disgard the max Point from the heap. Repeat this until all elements have
 * been added into the heap. The remaining k elements in the heap will be returned as the answer.
 * The addition and extraction operation of the heap takes O(lg(k)) time, and the total run time is O(n lg(k)).
 * As Java JDK provides no such priority queue container, I'll use a TreeSet to mock it.
 */

import java.util.TreeSet;


public class OriginQuestion {

  public static class Point  implements Comparable<Point> {

    public double x;

    public double y;

    // compareTo defines the order of Point in the MaxHeap
    public int compareTo(Point o) {
        double d1 = x * x + y * y;
    	double d2 = o.x * o.x + o.y * o.y;
    	
    	return Double.compare(d2, d1);
    }
    
    public String toString() {
    	return "(" + x + ", " + y + ")";
    }
  }

  public static Point[] closestk( Point  myList[], int k ) {
      
      if (null == myList || k <= 0 || k > myList.length) {
          return null;
      }
	  
      // Use TreeSet to imitate a heap
      TreeSet<Point> maxHeap = new TreeSet<Point>();
      int i = 0;

      // Build the heap of size k
      for (i = 0; i < k; i++) {
          maxHeap.add(myList[i]);
      }
      
      // Everytime add a Point into the heap, remove the largest one.
      for (i = k; i < myList.length; i++) {
          maxHeap.add(myList[i]);
          maxHeap.remove(maxHeap.first());
      }
      
      // The remaining elements are the answer.
      Point[] result = new Point[k];
      for (i = 0; i < k; i++) {
          result[i] = maxHeap.first();
          maxHeap.remove(maxHeap.first());
      }
      return result;
  }

  public static void main(String[] args) {
	  int n = 10;
	  int k = 5;
	  Point[] points = new Point[n];
	for (int i = 0; i < n; i++) {
		Point p = new Point();
		p.x = i;
		p.y = i;
		points[i] = p;
	}
	
	Point[] result = closestk(points, k);
	for (int i = 0; i < result.length; i++) {
		System.out.println(result[i]);
	}
  }
}
